import torch
import torch.nn as nn
import torchvision.models as models


class FCN8s(nn.Module):
    """
    基于resnet-50 骨干网络的 FCN-8s
    """

    def __init__(self, num_classes=2):
        super().__init__()

        resnet = models.resnet50(pretrained=True)

        # 下采样四倍
        self.conv_initial = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
        )

        self.layer1 = resnet.layer1 # 输出 特征长度不变
        self.layer2 = resnet.layer2 # 输出 下采样8倍
        self.layer3 = resnet.layer3 # 输出 下采样 16倍
        self.layer4 = resnet.layer4 # 输出下采样 32倍

        # 跳跃链接
        self.skip8_layer2 = nn.Conv2d(512, num_classes, kernel_size=1, stride=1, padding=0)
        self.skip16_layer3 = nn.Conv2d(1024, num_classes, kernel_size=1, stride=1, padding=0)
        self.skip32_layer4 = nn.Conv2d(2048, num_classes, kernel_size=1, stride=1, padding=0)

        # 2倍上采样模块 4-2-1 组合
        self.up_to_2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=4, stride=2, padding=1, bias=True)

        # 8倍上采样模块 16-8-4 组合
        self.up_to_8 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=16, stride=8, padding=4, bias=True)

    def forward(self, x):
        x = self.conv_initial(x)
        x = self.layer1(x)

        # 不同下采样长度，为了方便跳跃，所以先存储
        feat8 = self.layer2(x)
        feat16 = self.layer3(feat8)
        feat32 = self.layer4(feat16)

        # 使用1x1卷积，统一通道
        skip8 = self.skip8_layer2(feat8)
        skip16 = self.skip16_layer3(feat16)
        skip32 = self.skip32_layer4(feat32)

        # 上采样和跳转
        up_32_to_16 = self.up_to_2(skip32) # 此时 特征长度 1/16
        feat_blue = skip16 +  up_32_to_16

        feat_blue_up_to8 = self.up_to_2(feat_blue)
        feat_green = skip8 + feat_blue_up_to8

        feat_green_up_to_original = self.up_to_8(feat_green) # 原图尺寸

        return feat_green_up_to_original


if __name__ == "__main__":
    # # 通过查看renet50的结构，知晓对应下采样大小(尺寸长度)
    # model = models.resnet50(pretrained=True)
    #
    # print(model)

    model = FCN8s(num_classes=2)
    print(model)

    x = torch.randn(1, 3, 224, 224)
    # 运行前向传播
    out = model(x)
    print("模型输出尺寸:", out.shape)








